# get-pdb.py
# Download PDB file from 'www.rcsb.org' and store them in a series of names

import urllib2
import math

class atom:
    dic = { 'C' : 6, 'N' : 7, 'O' : 8, 'S' : 16 }
    def __init__(self, name, coordinate):
        self.name = name
        self.charge = self.dic[name]
        self.coordinate = coordinate

    def distance(self, other):
        dis_square = 0
        for i in range(3):
            dis_square += math.pow(self.coordinate[i] - other.coordinate[i], 2)
        
        dis = math.sqrt(dis_square)
        return dis

    def repulsion(self, other):
        if self.distance(other) != 0:
            return (self.charge*other.charge)/self.distance(other)
        else:
            return 0
            
url = "http://www.rcsb.org/pdb/download/downloadFile.do?fileFormat=pdb&compression=NO&structureId=1hlu"

file_name = '1HLU.xyz'
u = urllib2.urlopen(url)
f = open(file_name, 'w')
f.write(' \n')

# file_size_dl = 0
# block_size = 8192
num = 0
atom_info = []

while True:
    buffer = u.readline()
    if not buffer:
        break
    
    if buffer[0:6] == 'ATOM  ': 
        # save the coordinate info to a file
        element = buffer[77:78]
        x = buffer[31:38]
        y = buffer[39:46]
        z = buffer[47:54]    
        string = element + ' ' + x + ' ' + y + ' ' + z + '\n'
        f.write(string)
        
        # store the coordinate in a list
        coordinate = [ ]
        coordinate.append(float(x))
        coordinate.append(float(y))
        coordinate.append(float(z))
        
        atom_tmp = atom(element, coordinate)
        atom_info.append(atom_tmp)

    elif buffer[0:6] == 'MASTER':
        num = buffer[51:55]

# add the number of total atom to the top
f.seek(0,0)
f.write(num)
    
f.close()

# calculate the repulsion energy
repulsion = 0
for i in atom_info:
    for j in atom_info:
        repulsion += i.repulsion(j)

repulsion = 0.5*repulsion
print "The Repulsion Energy: ", repulsion
